package org.mockfu.junit;

import java.util.AbstractQueue;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Executes tests in parallel, allowing for subsets of tests to execute serially with respect to
 * members of the subset.
 */
public class ParallelExecutor
{
    private static final int DEFAULT_THREAD_COUNT = 5;

    private static volatile ThreadPoolExecutor executor;
    private static volatile SerializableQueue queue;


    @SuppressWarnings({"cast", "unchecked"})
    public static void initialize(int threads)
    {
        if (executor == null)
        {
            synchronized (ParallelSuite.class)
            {
                if (executor == null)
                {
                    if (threads <= 0)
                        threads = DEFAULT_THREAD_COUNT;
                    queue = new SerializableQueue();
                    executor = new ThreadPoolExecutor(threads, threads, 10l, TimeUnit.SECONDS, (BlockingQueue)queue);
                }
            }
        }

        if (executor.getMaximumPoolSize() < threads)
        {
            executor.setMaximumPoolSize(threads);
            executor.setCorePoolSize(threads);
        }
    }

    public static int threadCount()
    {
        return (executor == null) ? 1 : executor.getMaximumPoolSize();
    }

    /**
     * Submits a test to be executed in parallel. If {@code mutexNames} contains the empty string,
     * then no other tests will be able to execute concurrently with this test. Otherwise, for each
     * other value in {@code mutexNames} no test also submitted with that value will be able to
     * execute concurrently with this one.
     *
     * @param command
     * @param mutexNames
     * @return
     */
    public static Future<?> submit(Runnable command, String... mutexNames)
    {
        SerializableQueue.Entry entry = queue.new Entry(command, mutexNames);

        executor.execute(entry);

        return entry;
    }

    /**
     * A parallel task queue in which a subset of entries may be forced to execute serially.
     */
    private static class SerializableQueue extends AbstractQueue<SerializableQueue.Entry>
            implements BlockingQueue<SerializableQueue.Entry>
    {
        private final ReentrantLock lock = new ReentrantLock();
        private final Condition notEmpty = lock.newCondition();

        /* all guarded by lock */
        private final LinkedList<Entry> queue = new LinkedList<Entry>();
        private final Set<String> mutexes = new HashSet<String>();
        private int runningCount = 0;
        private boolean runningExclusive = false;


        /**
         * An entry in the queue and a wrapper around the task to be run. This is a little unusual
         * in that it's not the actual task submitted to the executor, it's added directly to the
         * queue, but it is the result returned to the caller of submit(). That's because the
         * executor may execute its task immediately, but we want to force it to always use the
         * queue.
         */
        private class Entry implements Future<Object>, Runnable
        {
            private final Runnable delegate;
            private final String[] mutexNames;

            private final Semaphore completionSignal = new Semaphore(0);
            private volatile boolean done = false;
            private volatile boolean canceled = false;
            private volatile boolean queued = false;
            private volatile Throwable exception;

            public Entry(Runnable delegate, String[] mutexNames)
            {
                this.delegate = delegate;
                this.mutexNames = mutexNames;
            }

            /* called with lock on monitor */
            public boolean isRunnable()
            {
                if (runningExclusive)
                    /* another task is running that requires total exclusion */
                    return false;

                if (mutexNames == null)
                    /* doesn't require exclusion */
                    return true;

                for (String mutexName : mutexNames)
                {
                    if ("".equals(mutexName))
                        /* must run in exclusion to everything else */
                        return runningCount == 0;

                    for (String mutex : mutexes)
                        if (mutex.equals(mutexName))
                            /* runs in exclusion to tasks with same mutex name */
                            return false;
                }
                return true;
            }

            /* called with lock on monitor */
            public void enableMutex()
            {
                if (mutexNames != null)
                {
                    for (String mutexName : mutexNames)
                    {
                        if ("".equals(mutexName))
                            runningExclusive = true;
                        else if (mutexName != null)
                            mutexes.add(mutexName);
                    }
                }

                runningCount++;
            }


            public void run()
            {
                if (!queued && !canceled)
                {
                    /* if executing as first task for a new thread, put entry on queue and do nothing */
                    offer(this);
                    return;
                }

                try
                {
                    if (canceled)
                        return;

                    delegate.run();
                }
                catch (Throwable ex)
                {
                    exception = ex;
                }
                finally
                {
                    /* remove the mutex */
                    lock.lock();
                    try
                    {
                        runningCount--;

                        if (mutexNames != null)
                        {
                            for (String mutexName : mutexNames)
                            {
                                if ("".equals(mutexName))
                                    runningExclusive = false;
                                else if (mutexName != null)
                                    mutexes.remove(mutexName);
                            }
                        }

                        /* mutex removal may allow new tasks to run */
                        notEmpty.signalAll();
                    }
                    finally
                    {
                        done = true;
                        /* signal to get() that the task is complete */
                        completionSignal.release();
                        lock.unlock();
                    }
                }
            }

            public boolean cancel(boolean mayInterruptIfRunning)
            {
                canceled = true;
                return done;
            }

            public Object get() throws InterruptedException, ExecutionException
            {
                completionSignal.acquire();
                if (exception != null)
                    throw new ExecutionException(exception);
                return null;
            }

            public Object get(long timeout, TimeUnit unit)
                    throws InterruptedException, ExecutionException, TimeoutException
            {
                if (!completionSignal.tryAcquire(timeout, unit))
                    throw new TimeoutException();
                if (exception != null)
                    throw new ExecutionException(exception);
                return null;
            }

            public boolean isCancelled()
            {
                return canceled;
            }

            public boolean isDone()
            {
                return done;
            }
        }


        @Override
        public Iterator<Entry> iterator()
        {
            throw new UnsupportedOperationException("Cannot iterate this queue");
        }

        @Override
        public int size()
        {
            lock.lock();
            try
            {
                int count = 0;

                for (Entry entry : queue)
                    if (entry.isRunnable())
                        count++;

                return count;
            }
            finally
            {
                lock.unlock();
            }
        }

        @Override
        public boolean isEmpty()
        {
            lock.lock();
            try
            {
                for (Entry entry : queue)
                    if (entry.isRunnable())
                        return false;

                return true;
            }
            finally
            {
                lock.unlock();
            }
        }

        public boolean offer(Entry o)
        {
            lock.lock();
            try
            {
                queue.offer(o);
                /* mark entry as executable */
                o.queued = true;
                notEmpty.signalAll();
                return true;
            }
            finally
            {
                lock.unlock();
            }
        }

        public boolean offer(Entry o, long timeout, TimeUnit unit)
        {
            return offer(o);
        }

        public void put(Entry o) throws InterruptedException
        {
            offer(o);
        }

        public Entry peek()
        {
            lock.lock();
            try
            {
                for (Entry entry : queue)
                    if (entry.isRunnable())
                        return entry;
                return null;
            }
            finally
            {
                lock.unlock();
            }
        }

        public Entry poll()
        {
            lock.lock();
            try
            {
                for (Iterator<Entry> i = queue.iterator(); i.hasNext(); )
                {
                    Entry entry = i.next();
                    if (entry.isRunnable())
                    {
                        i.remove();
                        entry.enableMutex();
                    }
                }

                return null;
            }
            finally
            {
                lock.unlock();
            }
        }

        public Entry poll(long timeout, TimeUnit unit) throws InterruptedException
        {
            long nanos = unit.toNanos(timeout);

            lock.lock();
            try
            {
                while (this.isEmpty())
                {
                    nanos = notEmpty.awaitNanos(nanos);
                    if (nanos <= 0)
                        return null;
                }

                for (Iterator<Entry> i = queue.iterator(); i.hasNext(); )
                {
                    Entry entry = i.next();
                    if (entry.isRunnable())
                    {
                        i.remove();
                        entry.enableMutex();
                        return entry;
                    }
                }

                throw new AssertionError("Unreachable code");
            }
            finally
            {
                lock.unlock();
            }
        }

        public Entry take() throws InterruptedException
        {
            lock.lock();
            try
            {
                while (this.isEmpty())
                    notEmpty.await();

                for (Iterator<Entry> i = queue.iterator(); i.hasNext(); )
                {
                    Entry entry = i.next();
                    if (entry.isRunnable())
                    {
                        i.remove();
                        entry.enableMutex();
                        return entry;
                    }
                }

                throw new AssertionError("Unreachable code");
            }
            finally
            {
                lock.unlock();
            }
        }

        public int drainTo(Collection<? super Entry> c)
        {
            return drainTo(c, Integer.MAX_VALUE);
        }

        public int drainTo(Collection<? super Entry> c, int maxElements)
        {
            if (maxElements <= 0)
                return maxElements;

            if (c == this)
                throw new IllegalArgumentException();

            lock.lock();
            try
            {
                int count = 0;
                for (Iterator<Entry> i = queue.iterator(); i.hasNext(); )
                {
                    Entry entry = i.next();
                    if (entry.isRunnable())
                    {
                        i.remove();
                        entry.enableMutex();
                        count++;
                    }
                }
                return count;
            }
            finally
            {
                lock.unlock();
            }
        }


        public int remainingCapacity()
        {
            return Integer.MAX_VALUE;
        }

    }

}
